package com.google.p2pbench;

import android.annotation.SuppressLint;
import android.net.nsd.NsdManager;
import android.net.nsd.NsdServiceInfo;
import android.os.Handler;
import android.os.HandlerThread;
import android.os.Looper;
import android.os.Message;
import android.os.Parcel;
import android.util.Log;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.net.SocketException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * A substitute for android.net.nsd.NsdManager which has caching issues.
 */
public class BetterNsdManager {
  private static final String TAG = "BetterNsdManager";
  private final Map<String, Set<NsdServiceInfo>> mRegisteredServices =
      Collections.synchronizedMap(new HashMap<String, Set<NsdServiceInfo>>());
  private final Map<NsdManager.RegistrationListener, NsdServiceInfo> mRegistrationListeners =
      Collections.synchronizedMap(new HashMap<NsdManager.RegistrationListener, NsdServiceInfo>());

  // TODO: TTL?
  private final Map<String, Map<String, NsdServiceInfo>> mDiscoveredServices =
      Collections.synchronizedMap(new HashMap<String, Map<String, NsdServiceInfo>>());

  private final Map<String, Set<NsdManager.DiscoveryListener>> mDiscoveries =
      Collections.synchronizedMap(new HashMap<String, Set<NsdManager.DiscoveryListener>>());
  private final Map<NsdManager.DiscoveryListener, String> mDiscoveryListeners =
      Collections.synchronizedMap(new HashMap<NsdManager.DiscoveryListener, String>());

  private static final int MSG_MAGIC = 0xc2fa6a4d;
  private static final byte MSG_TYPE_QUERY = 1;
  private static final byte MSG_TYPE_ANNOUNCE = 2;
  private static final int BROADCAST_PORT = 5354;
  private final SocketAddress mBroadcastAddress;
  private final DatagramSocket mSocket;

  private final HandlerThread mSenderThread = new HandlerThread("NsdSender");
  private final Sender mSender;
  @SuppressWarnings("unused")
  private final Receiver mReceiver;

  @SuppressLint("HandlerLeak")
  private class Sender extends Handler {
    Sender(Looper looper) {
      super(looper);
    }

    @Override
    public void handleMessage(Message msg) {
      DatagramPacket packet = (DatagramPacket) msg.obj;
      try {
        mSocket.send(packet);
      } catch (IOException e) {
        Log.w(TAG, "Failed to send packet", e);
      }
    }

    void queuePacket(DatagramPacket packet) {
      obtainMessage(0, packet).sendToTarget();
    }
  }

  private class Receiver implements Runnable {
    private final Thread mThread = new Thread(this, "NsdReceiver");

    private void handleQuery(String serviceType) {
      Set<NsdServiceInfo> infos = mRegisteredServices.get(serviceType);
      if (infos == null) return;
      synchronized (infos) {
        for (NsdServiceInfo info : infos) {
          mSender.queuePacket(createAnnouncement(info));
        }
      }
    }

    private void handleAnnounce(NsdServiceInfo info) {
      String serviceType = info.getServiceType();
      Map<String, NsdServiceInfo> services = mDiscoveredServices.get(serviceType);
      if (services == null) {
        services = Collections.synchronizedMap(new HashMap<String, NsdServiceInfo>(1));
        mDiscoveredServices.put(serviceType, services);
      }
      NsdServiceInfo previous = services.put(info.getServiceName(), info);
      if (info.equals(previous)) return; // Ignore duplicates.
      Set<NsdManager.DiscoveryListener> listeners = mDiscoveries.get(serviceType);
      if (listeners == null) return;
      synchronized (listeners) {
        for (NsdManager.DiscoveryListener listener : listeners) {
          listener.onServiceFound(info);
        }
      }
    }

    private void handlePacket(DatagramPacket packet) {
      Parcel parcel = Parcel.obtain();
      parcel.unmarshall(packet.getData(), packet.getOffset(), packet.getLength());
      parcel.setDataPosition(0); // WTF!?
      int magic = parcel.readInt();
      Log.w(TAG, "Receiver.handlePacket " + magic);
      if (magic != MSG_MAGIC) return;
      int type = parcel.readByte();
      Log.w(TAG, "Receiver got message " + type);
      switch (type) {
        case MSG_TYPE_QUERY:
          handleQuery(parcel.readString());
          break;
        case MSG_TYPE_ANNOUNCE:
          NsdServiceInfo info = parcel.readParcelable(NsdServiceInfo.class.getClassLoader());
          handleAnnounce(info);
          break;
        default:
      }
      parcel.recycle();
    }

    Receiver() {
      mThread.start();
    }

    @Override
    public void run() {
      byte[] buf = new byte[2000];
      DatagramPacket packet = new DatagramPacket(buf, buf.length);
      while (true) {
        try {
          mSocket.receive(packet);
        } catch (IOException e) {
          Log.w(TAG, "Failed to receive packet", e);
          return;
        }
        handlePacket(packet);
      }
    }
  }

  private DatagramPacket createQuery(String serviceType) {
    Parcel parcel = Parcel.obtain();
    parcel.writeInt(MSG_MAGIC);
    parcel.writeByte(MSG_TYPE_QUERY);
    parcel.writeString(serviceType);
    byte[] data = parcel.marshall();
    parcel.recycle();
    try {
      return new DatagramPacket(data, data.length, mBroadcastAddress);
    } catch (SocketException e) {
      Log.e(TAG, "WTF?", e);
    }
    return null;
  }

  private DatagramPacket createAnnouncement(NsdServiceInfo info) {
    Parcel parcel = Parcel.obtain();
    parcel.writeInt(MSG_MAGIC);
    parcel.writeByte(MSG_TYPE_ANNOUNCE);
    parcel.writeParcelable(info, 0);
    byte[] data = parcel.marshall();
    parcel.recycle();
    try {
      return new DatagramPacket(data, data.length, mBroadcastAddress);
    } catch (SocketException e) {
      Log.e(TAG, "WTF?", e);
    }
    return null;
  }

  public BetterNsdManager(InetAddress broadcastAddress) throws SocketException {
    mBroadcastAddress = new InetSocketAddress(broadcastAddress, BROADCAST_PORT);
    mSocket = new DatagramSocket(null);
    mSocket.setReuseAddress(true);
    mSocket.bind(new InetSocketAddress(BROADCAST_PORT));
    mSenderThread.start();
    mSender = new Sender(mSenderThread.getLooper());
    mReceiver = new Receiver();
  }

  public void shutdown() {
    mSocket.close();
  }

  public void registerService(
      NsdServiceInfo info, int protocolType, NsdManager.RegistrationListener listener) {

    Set<NsdServiceInfo> services = mRegisteredServices.get(info.getServiceType());
    if (services == null) {
      services = Collections.synchronizedSet(new HashSet<NsdServiceInfo>(1));
      mRegisteredServices.put(info.getServiceType(), services);
    }
    boolean added = services.add(info);
    assert added;
    mRegistrationListeners.put(listener, info);
    listener.onServiceRegistered(info);

    mSender.queuePacket(createAnnouncement(info));
  }

  public void unregisterService(NsdManager.RegistrationListener listener) {
    NsdServiceInfo info = mRegistrationListeners.remove(listener);
    Set<NsdServiceInfo> services = mRegisteredServices.get(info.getServiceType());
    services.remove(info);
    if (services.isEmpty()) mRegisteredServices.remove(info.getServiceType());
    listener.onServiceUnregistered(info);
  }

  public void discoverServices(
      String serviceType, int protocolType, NsdManager.DiscoveryListener listener) {

    Set<NsdManager.DiscoveryListener> listeners = mDiscoveries.get(serviceType);
    if (listeners == null) {
      listeners = Collections.synchronizedSet(new HashSet<NsdManager.DiscoveryListener>(1));
      mDiscoveries.put(serviceType, listeners);
    }
    boolean added = listeners.add(listener);
    assert added;
    mDiscoveryListeners.put(listener, serviceType);
    listener.onDiscoveryStarted(serviceType);

    mSender.queuePacket(createQuery(serviceType));

    Map<String, NsdServiceInfo> services = mDiscoveredServices.get(serviceType);
    if (services == null) return;

    List<NsdServiceInfo> infos = new ArrayList<NsdServiceInfo>(services.values());
    for (NsdServiceInfo info : infos) {
      listener.onServiceFound(info);
    }

    // TODO: keep sending queries for undiscovered discoveries?
  }

  public void stopServiceDiscovery(NsdManager.DiscoveryListener listener) {
    String serviceType = mDiscoveryListeners.remove(listener);
    Set<NsdManager.DiscoveryListener> listeners = mDiscoveries.get(serviceType);
    boolean removed = listeners.remove(listener);
    assert removed;
    if (listeners.isEmpty()) mDiscoveries.remove(serviceType);
    listener.onDiscoveryStopped(serviceType);
  }

  public void resolveService(NsdServiceInfo info, NsdManager.ResolveListener listener) {
    listener.onServiceResolved(info); // duh!
  }
}
